#!/usr/bin/env python3
import os
import cv2
import numpy as np
from dataclasses import dataclass
from typing import Optional, Tuple
import glob
import time  # 新增导入 time 模块
from shared_memory_reader import SharedMemoryReader
from utils import get_depth_image
from logger import setup_logger

logger = setup_logger()

@dataclass
class ImageData:
    """图像数据结构"""
    timestamp: str
    readable_timestamp: str
    color_img: np.ndarray
    depth_img: np.ndarray
    color_path: Optional[str] = None

class ImageProvider:
    """统一的图像提供接口"""
    
    def __init__(self, from_folder: bool = True, color_folder: str = "./data/Color1", depth_folder: str = "./data/Depth1"):
        self.from_folder = from_folder
        self.color_folder = color_folder
        self.depth_folder = depth_folder
        self.shm_reader = None
        self.color_files = []
        self.current_index = 0
        
        if from_folder:
            self.color_files = sorted(glob.glob(os.path.join(color_folder, "*.png")))
            logger.info(f"从文件夹读取，共找到 {len(self.color_files)} 个图片文件")
        else:
            self.shm_reader = SharedMemoryReader()
            logger.info("从共享内存读取")
    
    def get_next_image(self) -> Optional[ImageData]:
        """获取下一张图像"""
        if self.from_folder:
            return self._get_from_folder()
        else:
            return self._get_from_memory()
    
    def _get_from_folder(self) -> Optional[ImageData]:
        """从文件夹获取图像，按文件名中的数字大小排序"""
        # 按文件名中的数字排序
        self.color_files.sort(key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
        
        if self.current_index >= len(self.color_files):
            return None
        
        color_path = self.color_files[self.current_index]
        depth_path = os.path.join(self.depth_folder, os.path.basename(color_path))
        self.current_index += 1
        
        if not os.path.exists(depth_path):
            logger.warning(f"缺失深度图: {depth_path}")
            return self._get_from_folder()  # 递归获取下一张
        
        color_img = cv2.imread(color_path)
        depth_img = get_depth_image(color_path)
        
        if color_img is None or depth_img is None:
            logger.warning(f"读取图像失败: {color_path} 或 {depth_path}")
            return self._get_from_folder()  # 递归获取下一张
        
        timestamp = os.path.splitext(os.path.basename(color_path))[0]
        
        return ImageData(
            timestamp=timestamp,
            readable_timestamp=timestamp,
            color_img=color_img,
            depth_img=depth_img,
            color_path=color_path
        )

    def _get_from_memory(self) -> Optional[ImageData]:
        """从共享内存获取图像"""
        logger.info("开始从共享内存读取图像")
        start_time = time.time()  # 记录开始时间
        
        readable_timestamp, timestamp, color_img, depth_img = self.shm_reader.read_images()
        
        read_time = time.time() - start_time  # 计算读取耗时
        logger.info(f"从共享内存读取图像完成，耗时: {read_time:.4f} 秒")
        
        if color_img is None or depth_img is None:
            logger.warning("从共享内存读取图像失败")
            return None
        
        # 可选：保存图像到文件
        write_start_time = time.time()  # 记录写入开始时间
        success, color_path, depth_path = self.shm_reader.save_images(timestamp, color_img, depth_img)
        write_time = time.time() - write_start_time  # 计算写入耗时
        
        if success:
            logger.info(f"图像保存成功，路径: {color_path}, {depth_path}, 写入耗时: {write_time:.4f} 秒")
        else:
            logger.warning("图像保存失败")
        
        return ImageData(
            timestamp=timestamp,
            readable_timestamp=readable_timestamp,
            color_img=color_img,
            depth_img=depth_img,
            color_path=color_path if success else None
        )
    
    def close(self):
        """关闭资源"""
        if self.shm_reader:
            self.shm_reader.close()